load("//tensorflow:strict.default.bzl", "py_strict_library", "py_strict_test")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")

package(
    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
    default_visibility = [
        "//tensorflow:internal",
    ],
    licenses = ["notice"],
)

py_strict_library(
    name = "parallel_for",
    srcs = [
        "__init__.py",
    ],
    srcs_version = "PY3",
    deps = [
        ":control_flow_ops",
        ":gradients",
    ],
)

py_strict_library(
    name = "pfor_lib",
    srcs = ["pfor.py"],
    srcs_version = "PY3",
    deps = [
        "//tensorflow/compiler/tf2xla/python:xla",
        "//tensorflow/core:protos_all_py",
        "//tensorflow/python/eager:context",
        "//tensorflow/python/eager:def_function",
        "//tensorflow/python/eager:execute",
        "//tensorflow/python/framework:constant_op",
        "//tensorflow/python/framework:dtypes",
        "//tensorflow/python/framework:func_graph",
        "//tensorflow/python/framework:ops",
        "//tensorflow/python/framework:smart_cond",
        "//tensorflow/python/framework:sparse_tensor",
        "//tensorflow/python/framework:tensor",
        "//tensorflow/python/framework:tensor_shape",
        "//tensorflow/python/framework:tensor_util",
        "//tensorflow/python/ops:array_ops",
        "//tensorflow/python/ops:array_ops_gen",
        "//tensorflow/python/ops:array_ops_stack",
        "//tensorflow/python/ops:cond",
        "//tensorflow/python/ops:control_flow_assert",
        "//tensorflow/python/ops:control_flow_ops",
        "//tensorflow/python/ops:control_flow_switch_case",
        "//tensorflow/python/ops:data_flow_ops",
        "//tensorflow/python/ops:handle_data_util",
        "//tensorflow/python/ops:image_ops_gen",
        "//tensorflow/python/ops:linalg_ops",
        "//tensorflow/python/ops:linalg_ops_gen",
        "//tensorflow/python/ops:list_ops",
        "//tensorflow/python/ops:list_ops_gen",
        "//tensorflow/python/ops:manip_ops",
        "//tensorflow/python/ops:map_fn",
        "//tensorflow/python/ops:math_ops",
        "//tensorflow/python/ops:math_ops_gen",
        "//tensorflow/python/ops:nn_ops",
        "//tensorflow/python/ops:nn_ops_gen",
        "//tensorflow/python/ops:optional_ops_gen",
        "//tensorflow/python/ops:parsing_ops_gen",
        "//tensorflow/python/ops:random_ops_gen",
        "//tensorflow/python/ops:resource_variable_ops",
        "//tensorflow/python/ops:sparse_ops",
        "//tensorflow/python/ops:sparse_ops_gen",
        "//tensorflow/python/ops:special_math_ops",
        "//tensorflow/python/ops:spectral_ops_gen",
        "//tensorflow/python/ops:tensor_array_ops",
        "//tensorflow/python/ops:while_loop",
        "//tensorflow/python/platform:flags",
        "//tensorflow/python/platform:tf_logging",
        "//tensorflow/python/util:compat",
        "//tensorflow/python/util:nest",
        "//tensorflow/python/util:object_identity",
        "//third_party/py/numpy",
    ],
)

py_strict_library(
    name = "control_flow_ops",
    srcs = ["control_flow_ops.py"],
    srcs_version = "PY3",
    visibility = ["//visibility:public"],
    deps = [
        ":pfor_lib",
        "//tensorflow/python/autograph/core:ag_ctx",
        "//tensorflow/python/autograph/impl:api",
        "//tensorflow/python/eager:context",
        "//tensorflow/python/eager:def_function",
        "//tensorflow/python/framework:composite_tensor",
        "//tensorflow/python/framework:indexed_slices",
        "//tensorflow/python/framework:ops",
        "//tensorflow/python/framework:sparse_tensor",
        "//tensorflow/python/framework:tensor",
        "//tensorflow/python/framework:tensor_shape",
        "//tensorflow/python/framework:tensor_util",
        "//tensorflow/python/framework:type_spec",
        "//tensorflow/python/ops:array_ops",
        "//tensorflow/python/ops:cond",
        "//tensorflow/python/ops:math_ops",
        "//tensorflow/python/ops:tensor_array_ops",
        "//tensorflow/python/ops:while_loop",
        "//tensorflow/python/platform:tf_logging",
        "//tensorflow/python/util:nest",
        "//tensorflow/python/util:tf_decorator",
        "//tensorflow/python/util:tf_export",
        "//tensorflow/python/util:variable_utils",
    ],
)

py_strict_library(
    name = "test_util",
    srcs = ["test_util.py"],
    srcs_version = "PY3",
    deps = [
        ":control_flow_ops",
        "//tensorflow/python/ops:variables",
        "//tensorflow/python/platform:client_testlib",
        "//tensorflow/python/util:nest",
        "//third_party/py/numpy",
    ],
)

cuda_py_strict_test(
    name = "control_flow_ops_test",
    srcs = ["control_flow_ops_test.py"],
    shard_count = 16,
    tags = [
        "no_oss",
        "no_rocm",
    ],
    deps = [
        ":control_flow_ops",
        ":test_util",
        "//tensorflow/core:protos_all_py",
        "//tensorflow/python/client:session",
        "//tensorflow/python/eager:backprop",
        "//tensorflow/python/eager:context",
        "//tensorflow/python/eager:def_function",
        "//tensorflow/python/framework:composite_tensor",
        "//tensorflow/python/framework:config",
        "//tensorflow/python/framework:constant_op",
        "//tensorflow/python/framework:dtypes",
        "//tensorflow/python/framework:importer",
        "//tensorflow/python/framework:indexed_slices",
        "//tensorflow/python/framework:ops",
        "//tensorflow/python/framework:sparse_tensor",
        "//tensorflow/python/framework:tensor_shape",
        "//tensorflow/python/framework:tensor_spec",
        "//tensorflow/python/framework:test_lib",
        "//tensorflow/python/framework:type_spec",
        "//tensorflow/python/ops:array_ops",
        "//tensorflow/python/ops:bitwise_ops",
        "//tensorflow/python/ops:cond",
        "//tensorflow/python/ops:cond_v2",
        "//tensorflow/python/ops:control_flow_assert",
        "//tensorflow/python/ops:control_flow_v2_toggles",
        "//tensorflow/python/ops:data_flow_ops",
        "//tensorflow/python/ops:functional_ops",
        "//tensorflow/python/ops:gradient_checker_v2",
        "//tensorflow/python/ops:gradients",
        "//tensorflow/python/ops:image_ops",
        "//tensorflow/python/ops:list_ops",
        "//tensorflow/python/ops:list_ops_gen",
        "//tensorflow/python/ops:logging_ops",
        "//tensorflow/python/ops:manip_ops",
        "//tensorflow/python/ops:map_fn",
        "//tensorflow/python/ops:math_ops",
        "//tensorflow/python/ops:nn",
        "//tensorflow/python/ops:nn_ops_gen",
        "//tensorflow/python/ops:optional_ops_gen",
        "//tensorflow/python/ops:parsing_ops",
        "//tensorflow/python/ops:random_ops",
        "//tensorflow/python/ops:resource_variable_ops",
        "//tensorflow/python/ops:rnn_cell",
        "//tensorflow/python/ops:stateless_random_ops",
        "//tensorflow/python/ops:tensor_array_grad",
        "//tensorflow/python/ops:tensor_array_ops",
        "//tensorflow/python/ops:variables",
        "//tensorflow/python/ops:while_loop",
        "//tensorflow/python/ops/ragged:ragged_tensor",
        "//tensorflow/python/ops/signal:fft_ops",
        "//tensorflow/python/platform:client_testlib",
        "//tensorflow/python/util:nest",
        "//third_party/py/numpy",
        "@absl_py//absl/testing:parameterized",
    ],
)

cuda_py_strict_test(
    name = "xla_control_flow_ops_test",
    srcs = ["xla_control_flow_ops_test.py"],
    tags = [
        # XLA is not enabled by default on Mac or Windows.
        "no_mac",
        "no_windows",
        "no_gpu",  # TODO(b/155761551): Flaky on GPU on TAP
        "no_tfrt",  # Note: Legacy XLA test, which depends on EncapsulateXlaComputationsPass.
        "notsan",  # TODO(b/271486347): Fails TSAN on TAP
    ],
    xla_enabled = True,
    deps = [
        ":control_flow_ops",
        ":test_util",
        "//tensorflow/compiler/tf2xla/python:xla",
        "//tensorflow/python/compiler/xla",
        "//tensorflow/python/compiler/xla:compiler_py",
        "//tensorflow/python/eager:context",
        "//tensorflow/python/eager:def_function",
        "//tensorflow/python/framework:constant_op",
        "//tensorflow/python/framework:test_lib",
        "//tensorflow/python/ops:array_ops",
        "//tensorflow/python/ops:control_flow_v2_toggles",
        "//tensorflow/python/ops:math_ops",
        "//tensorflow/python/ops:random_ops",
        "//tensorflow/python/ops:resource_variable_ops",
        "//tensorflow/python/ops:while_loop",
        "//tensorflow/python/platform:client_testlib",
    ],
)

cuda_py_strict_test(
    name = "array_test",
    srcs = ["array_test.py"],
    tags = [
        "nogpu",  # b/217374776
    ],
    deps = [
        ":control_flow_ops",
        ":test_util",
        "//tensorflow/python/eager:backprop",
        "//tensorflow/python/framework:constant_op",
        "//tensorflow/python/framework:dtypes",
        "//tensorflow/python/framework:test_lib",
        "//tensorflow/python/ops:array_ops",
        "//tensorflow/python/ops:array_ops_stack",
        "//tensorflow/python/ops:math_ops",
        "//tensorflow/python/ops:nn",
        "//tensorflow/python/ops:random_ops",
        "//tensorflow/python/ops:tensor_array_grad",
        "//tensorflow/python/platform:client_testlib",
    ],
)

cuda_py_strict_test(
    name = "math_test",
    srcs = ["math_test.py"],
    shard_count = 5,
    tags = ["optonly"],  # Too slow in non-opt mode
    deps = [
        ":control_flow_ops",
        ":test_util",
        "//tensorflow/python/eager:backprop",
        "//tensorflow/python/framework:constant_op",
        "//tensorflow/python/framework:dtypes",
        "//tensorflow/python/framework:ops",
        "//tensorflow/python/framework:tensor",
        "//tensorflow/python/framework:test_lib",
        "//tensorflow/python/ops:array_ops",
        "//tensorflow/python/ops:clip_ops",
        "//tensorflow/python/ops:linalg_ops",
        "//tensorflow/python/ops:math_ops",
        "//tensorflow/python/ops:nn",
        "//tensorflow/python/ops:random_ops",
        "//tensorflow/python/ops:special_math_ops",
        "//tensorflow/python/ops:tensor_array_grad",
        "//tensorflow/python/platform:client_testlib",
        "@absl_py//absl/testing:parameterized",
    ],
)

py_strict_library(
    name = "gradients",
    srcs = ["gradients.py"],
    srcs_version = "PY3",
    deps = [
        ":control_flow_ops",
        "//tensorflow/python/framework:ops",
        "//tensorflow/python/framework:tensor",
        "//tensorflow/python/ops:array_ops",
        "//tensorflow/python/ops:check_ops",
        "//tensorflow/python/ops:gradients_impl",
        "//tensorflow/python/util:nest",
    ],
)

cuda_py_strict_test(
    name = "gradients_test",
    srcs = ["gradients_test.py"],
    tags = ["optonly"],  # Too slow in non-opt mode
    deps = [
        ":control_flow_ops",
        ":gradients",
        "//tensorflow/python/client:session",
        "//tensorflow/python/framework:constant_op",
        "//tensorflow/python/framework:dtypes",
        "//tensorflow/python/framework:errors",
        "//tensorflow/python/framework:ops",
        "//tensorflow/python/framework:test_lib",
        "//tensorflow/python/layers",
        "//tensorflow/python/ops:array_ops",
        "//tensorflow/python/ops:array_ops_stack",
        "//tensorflow/python/ops:functional_ops",
        "//tensorflow/python/ops:gradients",
        "//tensorflow/python/ops:math_ops",
        "//tensorflow/python/ops:nn",
        "//tensorflow/python/ops:random_ops",
        "//tensorflow/python/ops:rnn_cell",
        "//tensorflow/python/ops:variables",
        "//tensorflow/python/ops:while_loop",
        "//tensorflow/python/ops/losses",
        "//tensorflow/python/platform:client_testlib",
        "//tensorflow/python/util:nest",
        "//third_party/py/numpy",
    ],
)

py_strict_test(
    name = "pfor_test",
    srcs = ["pfor_test.py"],
    python_version = "PY3",
    srcs_version = "PY3",
    deps = [
        ":pfor_lib",
        "//tensorflow/python/framework:constant_op",
        "//tensorflow/python/framework:dtypes",
        "//tensorflow/python/framework:ops",
        "//tensorflow/python/framework:tensor",
        "//tensorflow/python/ops:array_ops",
        "//tensorflow/python/ops:random_ops",
        "//tensorflow/python/platform:client_testlib",
    ],
)
